"""
RSI策略(pyfolio实现)
策略逻辑：
1. 当RSI低于30时买入(超卖)
2. 当RSI高于70时卖出(超买)
"""

import pandas as pd
import pyfolio as pf
from empyrical import rsi

def rsi_strategy(prices, rsi_period=14, overbought=70, oversold=30):
    """
    基于RSI的交易策略
    
    参数:
        prices: 价格序列
        rsi_period: RSI计算周期
        overbought: 超买阈值
        oversold: 超卖阈值
    
    返回:
        positions: 持仓序列
    """
    rsi_values = rsi(prices, window=rsi_period)
    positions = pd.Series(0, index=prices.index)
    
    # 生成交易信号
    positions[rsi_values < oversold] = 1  # 买入信号
    positions[rsi_values > overbought] = -1  # 卖出信号
    
    return positions

if __name__ == '__main__':
    # 示例用法
    import yfinance as yf
    
    # 获取示例数据
    data = yf.download('AAPL', start='2020-01-01', end='2023-01-01')
    prices = data['Close']
    
    # 运行策略
    positions = rsi_strategy(prices)
    
    # 使用pyfolio分析结果
    returns = positions.shift(1) * prices.pct_change()
    pf.create_returns_tear_sheet(returns)